from __future__ import annotations

import sys
from io import BytesIO

import pytest

from PIL import Image

iterations = 5000


"""
When run on a system without the jpeg leak fixes,
the valgrind runs look like this.

valgrind --tool=massif python test-installed.py -s -v checks/check_jpeg_leaks.py

"""


pytestmark = pytest.mark.skipif(
    sys.platform.startswith("win32"), reason="requires Unix or macOS"
)

"""
pre patch:

    MB
31.62^                                                                       :
     |                                                              @:@:@:@#::
     |                                                     @:@:@@:@:@:@:@:@#::
     |                                             ::::::::@:@:@@:@:@:@:@:@#::
     |                                   :::::@::::::: ::::@:@:@@:@:@:@:@:@#::
     |                          @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |               ::::::@::::@:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |          ::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |         :::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |        ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |        ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |      ::::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |      : ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |     @: ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |    @@: ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |   :@@: ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     |   :@@: ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     | :@:@@: ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     | :@:@@: ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
     | :@:@@: ::::::::: : :@: : @:::::::::::::@: : ::: ::::@:@:@@:@:@:@:@:@#::
   0 +----------------------------------------------------------------------->Gi
     0                                                                   8.535


post-patch:

    MB
21.03^          :::@@:::@::::@@:::::::@@::::::::@::::::::::::@:::@:::::::@::::
     |         #:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |         #:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |        :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |        :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |        :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |      :::#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |      : :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |      : :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |     @: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |    @@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |    @@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |    @@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |   :@@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     |   :@@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     | :@:@@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     | :@:@@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     | :@:@@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     | :@:@@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
     | :@:@@: :#:::@ :::@::::@ : :: : @ :::::: :@:: ::: :::: @:: @:::::::@::::
   0 +----------------------------------------------------------------------->Gi
     0                                                                   8.421

"""


standard_l_qtable = (
    # fmt: off
    16, 11, 10, 16,  24,  40,  51,  61,
    12, 12, 14, 19,  26,  58,  60,  55,
    14, 13, 16, 24,  40,  57,  69,  56,
    14, 17, 22, 29,  51,  87,  80,  62,
    18, 22, 37, 56,  68, 109, 103,  77,
    24, 35, 55, 64,  81, 104, 113,  92,
    49, 64, 78, 87, 103, 121, 120, 101,
    72, 92, 95, 98, 112, 100, 103,  99,
    # fmt: on
)

standard_chrominance_qtable = (
    # fmt: off
    17, 18, 24, 47, 99, 99, 99, 99,
    18, 21, 26, 66, 99, 99, 99, 99,
    24, 26, 56, 99, 99, 99, 99, 99,
    47, 66, 99, 99, 99, 99, 99, 99,
    99, 99, 99, 99, 99, 99, 99, 99,
    99, 99, 99, 99, 99, 99, 99, 99,
    99, 99, 99, 99, 99, 99, 99, 99,
    99, 99, 99, 99, 99, 99, 99, 99,
    # fmt: on
)


@pytest.mark.parametrize(
    "qtables",
    (
        (standard_l_qtable, standard_chrominance_qtable),
        [standard_l_qtable, standard_chrominance_qtable],
    ),
)
def test_qtables_leak(qtables: tuple[tuple[int, ...]] | list[tuple[int, ...]]) -> None:
    with Image.open("Tests/images/hopper.ppm") as im:
        for _ in range(iterations):
            test_output = BytesIO()
            im.save(test_output, "JPEG", qtables=qtables)


def test_exif_leak() -> None:
    """
    pre patch:

        MB
    177.1^                                                                       #
         |                                                                    @@@#
         |                                                                :@@@@@@#
         |                                                             ::::@@@@@@#
         |                                                         ::::::::@@@@@@#
         |                                                     @@::::: ::::@@@@@@#
         |                                                  @@@@ ::::: ::::@@@@@@#
         |                                               @@@@@@@ ::::: ::::@@@@@@#
         |                                           @@::@@@@@@@ ::::: ::::@@@@@@#
         |                                        @@@@ : @@@@@@@ ::::: ::::@@@@@@#
         |                                   @@@@@@ @@ : @@@@@@@ ::::: ::::@@@@@@#
         |                                @@@@ @@ @ @@ : @@@@@@@ ::::: ::::@@@@@@#
         |                            @::@@ @@ @@ @ @@ : @@@@@@@ ::::: ::::@@@@@@#
         |                        ::::@: @@ @@ @@ @ @@ : @@@@@@@ ::::: ::::@@@@@@#
         |                     :@@: : @: @@ @@ @@ @ @@ : @@@@@@@ ::::: ::::@@@@@@#
         |                ::@@::@@: : @: @@ @@ @@ @ @@ : @@@@@@@ ::::: ::::@@@@@@#
         |            @@::: @ ::@@: : @: @@ @@ @@ @ @@ : @@@@@@@ ::::: ::::@@@@@@#
         |         @::@ : : @ ::@@: : @: @@ @@ @@ @ @@ : @@@@@@@ ::::: ::::@@@@@@#
         |      :::@: @ : : @ ::@@: : @: @@ @@ @@ @ @@ : @@@@@@@ ::::: ::::@@@@@@#
         |   @@@:: @: @ : : @ ::@@: : @: @@ @@ @@ @ @@ : @@@@@@@ ::::: ::::@@@@@@#
       0 +----------------------------------------------------------------------->Gi
         0                                                                   11.37


    post patch:

        MB
    21.06^        ::::::::::::::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |      ##::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |      # ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |      # ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |      # ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |     @# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |     @# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |     @# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |     @# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |    @@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |   @@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |   @@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |   @@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |   @@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         |   @@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         | @@@@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         | @ @@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         | @ @@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         | @ @@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
         | @ @@@# ::: ::::: : ::::::::::@::::@::::@::::@::::@::::@:::::::::@::::::
       0 +----------------------------------------------------------------------->Gi
         0                                                                   11.33
    """
    exif = b"12345678" * 4096

    with Image.open("Tests/images/hopper.ppm") as im:
        for _ in range(iterations):
            test_output = BytesIO()
            im.save(test_output, "JPEG", exif=exif)


def test_base_save() -> None:
    """
    base case:
        MB
    20.99^           :::::         :::::::::::::::::::::::::::::::::::::::::::@:::
         |         ##: : ::::::@::::::: :::: :::: : : : : : : :::::::::::: :::@:::
         |         # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |         # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |         # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |       @@# : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |       @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |       @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |     @@@ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |     @ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |    @@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |    @@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |    @@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |    @@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         |    @@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         | :@@@@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         | :@ @@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         | :@ @@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         | :@ @@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
         | :@ @@ @ # : : :: :: @:: :::: :::: :::: : : : : : : :::::::::::: :::@:::
       0 +----------------------------------------------------------------------->Gi
         0                                                                   7.882"""
    with Image.open("Tests/images/hopper.ppm") as im:
        for _ in range(iterations):
            test_output = BytesIO()
            im.save(test_output, "JPEG")
